"""GTest 扩展处理器

提供 GTest 单元测试的装饰器和扩展处理功能。
"""

from typing import cast, Optional
from mcpy.ir import IRContext, IRFunctionNode, IRModule, IRNode
from mcpy.extension_base import extension_processor
from mcpy.ir.expressions import IRCall, IRName, IRAttribute
from mcpy.ir.statements import IRStmtBlock, IRExprStmt
from mcpy.ir.declarations.import_node import IRImportNode, ImportKind
from .ir_nodes import IRGTestFunctionNode, IRGTestFixtureFunctionNode


def _ensure_test_import(module: IRModule, context=None) -> None:
    """确保模块中有 Test 的导入（从 gtest.testing 导入）

    检查模块中是否已经有 `from gtest.testing import Test` 的导入，
    如果没有则创建一个并添加到模块的开头。

    Args:
        module: IR 模块
        context: IR 上下文（可选，用于获取 compiler）
    """
    # 检查是否已经有 Test 的导入
    for node in module.nodes:
        if (
            isinstance(node, IRImportNode)
            and node.import_kind == ImportKind.FROM_IMPORT
            and node.source_module == "gtest.testing"
            and "Test" in node.imported_symbols
        ):
            # 已经存在 Test 的导入，如果未解析则解析它
            if not node.is_resolved:
                compiler = (
                    context.compiler
                    if context and context.compiler
                    else node._get_compiler()
                )
                if compiler:
                    node.resolve_import(compiler)
            return

    # 创建新的导入节点
    import_node = IRImportNode(
        source_line=0,  # 导入语句通常在文件开头
        parent=module,
        import_kind=ImportKind.FROM_IMPORT,
        source_module="gtest.testing",
        imported_symbols=["Test"],
    )

    # 将导入节点添加到模块的开头（在第一个节点之前）
    if len(module.nodes) > 0:
        module.insert_node_before(module.nodes[0], import_node)
    else:
        module.add_node(import_node)

    # 立即解析导入节点
    compiler = (
        context.compiler
        if context and context.compiler
        else import_node._get_compiler()
    )
    if compiler:
        import_node.resolve_import(compiler)


def _ensure_ensure_future_import(module: IRModule, context=None) -> None:
    """确保模块中有 ensure_future 的导入

    检查模块中是否已经有 `from mc.future import ensure_future` 的导入，
    如果没有则创建一个并添加到模块的开头。
    如果已经有 `from mc.future import ...` 的导入但没有 ensure_future，
    则添加到现有的导入中。

    创建或更新导入节点后，会自动调用 resolve_import 解析导入，
    resolve_import 会从模块树向上查找 compiler 并完成符号导入。

    Args:
        module: IR 模块
        context: IR 上下文（可选，用于获取 compiler）
    """
    # 检查是否已经有 ensure_future 的导入，或者是否有 from mc.future import ... 的导入
    existing_import: Optional[IRImportNode] = None
    for node in module.nodes:
        if (
            isinstance(node, IRImportNode)
            and node.import_kind == ImportKind.FROM_IMPORT
            and node.source_module == "mc.future"
        ):
            if "ensure_future" in node.imported_symbols:
                # 已经存在 ensure_future 的导入，如果未解析则解析它
                if not node.is_resolved:
                    # 尝试从 context 获取 compiler，否则从模块树向上查找
                    compiler = (
                        context.compiler
                        if context and context.compiler
                        else node._get_compiler()
                    )
                    if compiler:
                        node.resolve_import(compiler)
                return
            else:
                # 找到了 from mc.future import ... 的导入，但没有 ensure_future
                existing_import = node
                break

    if existing_import:
        # 添加到现有的导入中
        if "ensure_future" not in existing_import.imported_symbols:
            existing_import.imported_symbols.append("ensure_future")
        # 如果导入节点还未解析，立即解析它
        if not existing_import.is_resolved:
            # 尝试从 context 获取 compiler，否则从模块树向上查找
            compiler = (
                context.compiler
                if context and context.compiler
                else existing_import._get_compiler()
            )
            if compiler:
                existing_import.resolve_import(compiler)
    else:
        # 创建新的导入节点
        import_node = IRImportNode(
            source_line=0,  # 导入语句通常在文件开头
            parent=module,
            import_kind=ImportKind.FROM_IMPORT,
            source_module="mc.future",
            imported_symbols=["ensure_future"],
        )

        # 将导入节点添加到模块的开头（在第一个节点之前）
        if len(module.nodes) > 0:
            module.insert_node_before(module.nodes[0], import_node)
        else:
            module.add_node(import_node)

        # 立即解析导入节点，尝试从 context 获取 compiler，否则从模块树向上查找
        compiler = (
            context.compiler
            if context and context.compiler
            else import_node._get_compiler()
        )
        if compiler:
            import_node.resolve_import(compiler)


def _wrap_async_test_function(
    async_func: IRFunctionNode, module: IRModule, wrapper_name: str, context=None
) -> IRFunctionNode:
    """为异步测试函数创建同步包装函数

    将异步函数包装为同步函数，在包装函数中使用 ensure_future().get() 阻塞调用。

    Args:
        async_func: 异步函数节点
        module: IR 模块
        wrapper_name: 包装函数的名称（通常是原函数名）
        context: IR 上下文（可选，用于获取 compiler）

    Returns:
        同步包装函数节点
    """
    # 0. 确保模块中有 ensure_future 的导入
    _ensure_ensure_future_import(module, context)

    # 1. 重命名原异步函数为内部函数（添加 _async_ 前缀）
    internal_name = f"_async_{async_func.name}"
    async_func.name = internal_name

    # 2. 创建同步包装函数
    # 包装函数使用原函数名，但是是同步的
    wrapper_func = IRFunctionNode(
        name=wrapper_name,
        source_line=async_func.source_line,
        decorators=async_func.decorators,
        doc=async_func.doc,
        params=async_func.params,
        return_type=async_func.return_type,
        body=IRStmtBlock(parent=None),
        is_async=False,  # 包装函数是同步的
        is_static=async_func.is_static,
        extension_params=async_func.extension_params.copy(),
        parent=async_func.parent,
    )
    # 设置 body 的 parent，确保符号查找能正确向上遍历
    wrapper_func.body.parent = wrapper_func

    # 3. 构建包装函数的 body：ensure_future(_async_func(...)).get()
    # 3.1 创建内部函数调用：_async_func(...)
    # 传递所有函数参数
    internal_call = IRCall()
    internal_call.func = IRName(id=internal_name)
    internal_call.func.parent = internal_call
    # 为每个参数创建 IRName 引用
    internal_call.args = []
    for param in async_func.params:
        param_ref = IRName(id=param.name)
        param_ref.parent = internal_call
        internal_call.args.append(param_ref)
    internal_call.parent = wrapper_func.body

    # 3.1.1 为内部函数调用设置扩展参数，标记为全局异步函数调用
    # 这样 CallEmitter 能够正确识别并生成 struct_name{}(args) 形式的调用
    from mcpy.codegen.emitters.nested_function.utils import get_functor_name

    ir_module = async_func.get_ir_module()
    if ir_module:
        struct_name = get_functor_name(async_func)
        internal_call.set_extension_param("_global_async_struct_name", struct_name)

    # 3.2 创建 ensure_future 调用：ensure_future(_async_func())
    ensure_future_call = IRCall()
    ensure_future_call.func = IRName(id="ensure_future")
    ensure_future_call.func.parent = ensure_future_call
    ensure_future_call.args = [internal_call]
    internal_call.parent = ensure_future_call
    ensure_future_call.parent = wrapper_func.body

    # 3.3 创建 .get() 调用：ensure_future(_async_func()).get()
    get_attr = IRAttribute()
    get_attr.value = ensure_future_call
    get_attr.value.parent = get_attr
    get_attr.attr = "get"
    get_attr.parent = wrapper_func.body

    get_call = IRCall()
    get_call.func = get_attr
    get_call.func.parent = get_call
    get_call.args = []
    get_call.parent = wrapper_func.body

    # 3.4 创建表达式语句
    expr_stmt = IRExprStmt()
    expr_stmt.expr = get_call
    get_call.parent = expr_stmt
    expr_stmt.parent = wrapper_func.body
    expr_stmt.source_line = async_func.source_line

    # 设置 body 的语句列表
    wrapper_func.body.statements = [expr_stmt]
    return wrapper_func


@extension_processor(
    "gtest_test",
    match_rule=[{"decorator": "test", "node_type": "function"}],
    priority=50,
    stage="pre_analysis",
    timing="pre",
)
def process_gtest_test(context: IRContext):
    """处理 @test 装饰器，生成 GTest 测试用例"""
    func_node = cast(IRFunctionNode, context.node)

    # 获取装饰器参数
    decorator = context.get_decorator("test")
    assert decorator, "test 装饰器处理失败"

    # 获取测试名称，默认使用函数名
    test_name = func_node.name
    test_case_name = "DefaultTestSuite"

    if "args" in decorator and len(decorator["args"]) > 0:
        test_case_name = decorator["args"][0]
        if len(decorator["args"]) > 1:
            test_name = decorator["args"][1]

    # 如果函数是异步的，需要创建包装函数
    if func_node.is_async:
        # 创建同步包装函数
        # 注意：_wrap_async_test_function 会将原异步函数重命名为 _async_xxx
        # 原异步函数仍然在模块中，需要保留它以便生成
        wrapper_func = _wrap_async_test_function(
            func_node, context.module, func_node.name, context
        )
        # 在原异步函数位置后插入包装函数，保持源码顺序
        context.module.insert_node_after(func_node, wrapper_func)
        # 使用包装函数创建 GTest 节点
        gtest_node = IRGTestFunctionNode(
            name=wrapper_func.name,
            source_line=wrapper_func.source_line,
            decorators=wrapper_func.decorators,
            doc=wrapper_func.doc,
            params=wrapper_func.params,
            return_type=wrapper_func.return_type,
            body=wrapper_func.body,
            is_async=False,  # 包装函数是同步的
            is_static=wrapper_func.is_static,
            macro_type="TEST",
            test_suite=test_case_name,
            test_name=test_name,
            extension_params=wrapper_func.extension_params,
        )
        # 替换包装函数节点（原异步函数保留在模块中，不会被替换）
        context.module.replace_node(wrapper_func, gtest_node)
        # 注意：不要替换原异步函数，它已经被重命名为 _async_xxx，需要保留以便生成
        return
    else:
        # 同步函数，直接创建 GTest 节点
        gtest_node = IRGTestFunctionNode(
            name=func_node.name,
            source_line=func_node.source_line,
            decorators=func_node.decorators,
            doc=func_node.doc,
            params=func_node.params,
            return_type=func_node.return_type,
            body=func_node.body,
            is_async=func_node.is_async,
            is_static=func_node.is_static,
            macro_type="TEST",
            test_suite=test_case_name,
            test_name=test_name,
            extension_params=func_node.extension_params,
        )
        # 替换原始节点
        context.module.replace_node(func_node, gtest_node)

    # 添加必要的头文件
    context.module.require_include("<gtest/gtest.h>")


@extension_processor(
    "gtest_fixture",
    match_rule=[{"decorator": "fixture", "node_type": "class"}],
    priority=50,
    stage="pre_analysis",
    timing="pre",
)
def process_gtest_fixture(context: IRContext):
    """处理 @fixture 装饰器，生成 GTest 测试夹具"""

    if context.node.node_type != "class":
        return

    from mcpy.ir import IRClassNode

    class_node = cast(IRClassNode, context.node)

    # 获取装饰器参数
    decorator = context.get_decorator("fixture")
    assert decorator, "fixture 装饰器处理失败"

    # 确保导入了 Test 类（从 gtest.testing 导入）
    # 这样类型推导系统才能找到 Test 符号
    _ensure_test_import(context.module, context)

    # 设置基类为 Test（Python 侧的类型）
    # 代码生成时会根据 Test 的来源模块（gtest.testing）自动添加命名空间前缀
    class_node.bases = ["Test"]

    # 添加必要的头文件
    context.module.require_include("<gtest/gtest.h>")


@extension_processor(
    "gtest_fixture_test",
    match_rule=[{"decorator": "fixture", "node_type": "function"}],
    priority=50,
    stage="pre_analysis",
    timing="pre",
)
def process_gtest_fixture_test(context: IRContext):
    """处理 @fixture 装饰器，生成基于夹具的 GTest 测试用例"""

    if context.node.node_type != "function":
        return

    func_node = cast(IRFunctionNode, context.node)

    # 获取装饰器参数
    decorator = context.get_decorator("fixture")
    assert decorator, "fixture 装饰器处理失败"

    # 获取测试名称和夹具类名
    test_name = func_node.name
    fixture_class = "DefaultFixture"

    if "args" in decorator and len(decorator["args"]) > 0:
        fixture_class = decorator["args"][0]
        if len(decorator["args"]) > 1:
            test_name = decorator["args"][1]

    # 如果函数是异步的，需要创建包装函数
    if func_node.is_async:
        # 创建同步包装函数
        wrapper_func = _wrap_async_test_function(
            func_node, context.module, func_node.name, context
        )
        # 在原异步函数位置后插入包装函数，保持源码顺序
        # 原异步函数已经重命名为 _async_xxx，但仍然在模块中的原位置
        context.module.insert_node_after(func_node, wrapper_func)
        # 使用包装函数创建 GTest 节点
        gtest_node = IRGTestFixtureFunctionNode(
            name=wrapper_func.name,
            source_line=wrapper_func.source_line,
            decorators=wrapper_func.decorators,
            doc=wrapper_func.doc,
            params=wrapper_func.params,
            return_type=wrapper_func.return_type,
            body=wrapper_func.body,
            is_async=False,  # 包装函数是同步的
            is_static=wrapper_func.is_static,
            macro_type="TEST_F",
            test_suite=fixture_class,
            test_name=test_name,
            extension_params=wrapper_func.extension_params,
        )
        # 替换包装函数节点
        context.module.replace_node(wrapper_func, gtest_node)
    else:
        # 同步函数，直接创建 GTest 节点
        gtest_node = IRGTestFixtureFunctionNode(
            name=func_node.name,
            source_line=func_node.source_line,
            decorators=func_node.decorators,
            doc=func_node.doc,
            params=func_node.params,
            return_type=func_node.return_type,
            body=func_node.body,
            is_async=func_node.is_async,
            is_static=func_node.is_static,
            macro_type="TEST_F",
            test_suite=fixture_class,
            test_name=test_name,
            extension_params=func_node.extension_params,
        )
    # 替换原始节点
    context.module.replace_node(func_node, gtest_node)

    # 添加必要的头文件
    context.module.require_include("<gtest/gtest.h>")


@extension_processor(
    "gtest_test_f",
    match_rule=[{"decorator": "test_f", "node_type": "function"}],
    priority=50,
    stage="pre_analysis",
    timing="pre",
)
def process_gtest_test_f(context: IRContext):
    """处理 @test_f 装饰器，生成基于夹具的 GTest 测试用例"""

    if context.node.node_type != "function":
        return

    func_node = cast(IRFunctionNode, context.node)

    # 获取装饰器参数
    decorator = context.get_decorator("test_f")
    if not decorator:
        return

    # 获取测试名称和夹具类名
    test_name = func_node.name
    fixture_class = "DefaultFixture"

    if "args" in decorator and len(decorator["args"]) > 0:
        fixture_class = decorator["args"][0]
        if len(decorator["args"]) > 1:
            test_name = decorator["args"][1]

    # 如果函数是异步的，需要创建包装函数
    if func_node.is_async:
        # 创建同步包装函数
        wrapper_func = _wrap_async_test_function(
            func_node, context.module, func_node.name, context
        )
        # 在原异步函数位置后插入包装函数，保持源码顺序
        # 原异步函数已经重命名为 _async_xxx，但仍然在模块中的原位置
        context.module.insert_node_after(func_node, wrapper_func)
        # 使用包装函数创建 GTest 节点
        gtest_node = IRGTestFixtureFunctionNode(
            name=wrapper_func.name,
            source_line=wrapper_func.source_line,
            decorators=wrapper_func.decorators,
            doc=wrapper_func.doc,
            params=wrapper_func.params,
            return_type=wrapper_func.return_type,
            body=wrapper_func.body,
            is_async=False,  # 包装函数是同步的
            is_static=wrapper_func.is_static,
            macro_type="TEST_F",
            test_suite=fixture_class,
            test_name=test_name,
            extension_params=wrapper_func.extension_params,
        )
        # 替换包装函数节点
        context.module.replace_node(wrapper_func, gtest_node)
    else:
        # 同步函数，直接创建 GTest 节点
        gtest_node = IRGTestFixtureFunctionNode(
            name=func_node.name,
            source_line=func_node.source_line,
            decorators=func_node.decorators,
            doc=func_node.doc,
            params=func_node.params,
            return_type=func_node.return_type,
            body=func_node.body,
            is_async=func_node.is_async,
            is_static=func_node.is_static,
            macro_type="TEST_F",
            test_suite=fixture_class,
            test_name=test_name,
            extension_params=func_node.extension_params,
        )
    # 替换原始节点
    context.module.replace_node(func_node, gtest_node)

    # 添加必要的头文件
    context.module.require_include("<gtest/gtest.h>")
